import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, optimizers, losses, metrics
import matplotlib.pyplot as plt
import os
import cv2 as cv
import numpy as np

IMG_SIZE = 224
BASE_DIR, FILE_NAME = os.path.split(__file__)
dir = '../../../../large_data/CV4/_many_files/new_Gesture_Recognition'
IMG_DIR = os.path.join(BASE_DIR, dir)


class GestureRecognition(keras.Model):

    def __init__(self, n_cls, **kwargs):
        super().__init__(**kwargs)
        self.base_model = keras.applications.MobileNet(
            input_shape=(IMG_SIZE, IMG_SIZE, 3),
            include_top=False,
            weights='imagenet',
            pooling='avg',
        )
        self.base_model.trainable = False
        self.customer_model = layers.Dense(n_cls, activation=activations.softmax)

    def call(self, inputs, training=None, mask=None):
        x = self.base_model(inputs)
        x = self.customer_model(x)
        return x


if '__main__' == __name__:

    model = GestureRecognition(3)
    model.build(input_shape=(None, *(list(x.shape[1:]))))
    model.summary()

    x, y = [], []
    idx2label, label2idx = {}, {}
    yi = 0
    cnt = 0
    for sub_dir_name in os.listdir(IMG_DIR):
        sub_dir_path = os.path.join(IMG_DIR, sub_dir_name)
        print(f'Loading {sub_dir_name} ...')
        idx2label[yi] = sub_dir_name
        label2idx[sub_dir_name] = yi
        for file_name in os.listdir(sub_dir_path):
            file_path = os.path.join(sub_dir_path, file_name)
            img = cv.imread(file_path, cv.IMREAD_COLOR)
            img = cv.resize(img, (IMG_SIZE, IMG_SIZE))
            img = np.float32(img) / 255.
            x.append(img)
            y.append(yi)
            cnt += 1
            if cnt % 25 == 0:
                print(f'Loaded {cnt} pictures.')
        yi += 1
    x = np.float32(x)
    y = np.int64(y)

